package nachos.vm;

import nachos.machine.*;
import nachos.threads.*;
import nachos.userprog.*;
import nachos.vm.*;

import java.util.Random;
import java.util.Map;
import java.util.HashMap;
import java.util.List;
import java.util.Arrays;
import java.util.Formatter;

import static nachos.machine.Lib.debug;

/**
 * A <tt>UserProcess</tt> that supports demand-paging.
 */
public class VMProcess extends UserProcess {

    public static int numPageFaults = 0;
    public static int numPageAccesses = 0;

    // TODO: Let the swap grow as necessary
    private static final int SWAP_SIZE = 32;
    private static final String SWAP_FILENAME = "swap";

    private static final Lock swapLock = new Lock();

    /**
     * Allocate a new process.
     */
    public VMProcess() {
	super();
	if (ipt == null) {
	    ipt = new InvertedPageTable(Machine.processor().getNumPhysPages());
	    swap = new Swap(SWAP_SIZE, SWAP_FILENAME);


	}
	
	// this is for blanking out a physical memory space when a new stack/arg page is created
	if ( blankpage == null )
	    {
		blankpage = new byte[ Machine.processor().pageSize];
		for ( int i = 0; i < Machine.processor().pageSize; ++i )
		    {
			blankpage[i] = 0;
		    }
	    }
	
	pagelocations = new HashMap<Integer, Integer>(Machine.processor().getNumPhysPages());
	vpnToS  = new HashMap< Integer, Integer >( 25 ); // this number is arb.
	vpnToSI = new HashMap< Integer, Integer >( 25 ); // this number is arb.
    }

    /**
     * Pick a page resident in memory and swap it out to the
     * filesystem. Returns the physical page number that was freed.
     */
    public int swapOut() {
	// Find a page to evict
	int ppn = ipt.nextVictim();

	TranslationEntry vmEntry = ipt.table[ppn].te;
	int thePid               = ipt.table[ppn].pid;
	
	// Write the page to swap, invalidate the vm entry, and free
	// the page
	if (vmEntry.dirty) {
	    swap.write(thePid, vmEntry.vpn, ppn, vmEntry.readOnly);
	}
	vmEntry.valid = false;
	UserKernel.freePage(ppn);

	// At this point getNextPage should return the page number
	// that we just freed.
	int gotPpn = UserKernel.getNextPage();
	debug(dbgH, "gotPpn = <" + gotPpn + "> and ppn was <" + ppn + ">" );
	if (gotPpn != ppn) {
	    throw new RuntimeException("The wrong page was freed!");
	}
	

	String msg = String.format("After swapping out vpn %d for pid " +
				   "%d from ppn %d, in, vm is %s and " +
				   "swap is %s",
				   vmEntry.vpn, thePid, ppn, ipt, swap);
	debug(dbgVM, msg);
	return ppn;

    }

    /**
     * Read the page specified by vpn for the current pid into memory
     * at the specified physical page number. This assumes that the
     * physical page is available.
     */
    public void swapIn(int vpn, int ppn) {

	boolean readonly = false;
	// Read the bytes from the swap file into physical memory.
	debug(dbgH, "swapIn - pid is <" + pid + ">" );
	Swap.Entry swapEntry = swap.read(pid, vpn, ppn);

	// if the swap didn't find it, we need to look for it in the local "full" page table
	if ( swapEntry == null )
	    {		
		debug(dbgH, "swapIn - swap didn't have the page. = " + pagelocations.get( vpn ) );
		int loc = pagelocations.get( vpn );
		switch ( loc )
		    {
		    case PL_SWAP:
		    case PL_MEMORY:
			// If the page is supposed to be in swap or
			//memory but it's not, it must have been a
			//readonly or non-dirty page that was wiped
			//from memory and not written to swap. So we
			//need to pull it back out of the swap file.
			//System.out.println("Error - a page was
			//thought to be in swap or memory but
			//wasn't");
		    case PL_COFF:
			int s, i;
			s = vpnToS.get( vpn );
			i = vpnToSI.get( vpn );
			coff.getSection( s ).loadPage( i, ppn );
			readonly = coff.getSection(s).isReadOnly();
			break;

		    case PL_ST_ARG:
			debug(dbgH, "swapIn - found a stack or argument section to create.");
			// need to "create" the page data of zeroes
			// can simple zero out the stack page
			int realaddr = Processor.makeAddress( ppn,  0 );			
			System.arraycopy( blankpage, 
					  0, 
					  Machine.processor().getMemory(),
					  realaddr, 
					  Machine.processor().pageSize );
			readonly = false;
			break;

		    case PL_UNKNOWN:
		    default:
			System.out.println("Error - couldn't find a page in the locations. <" + vpn + ">" );
			terminateProcess( -1 );
		    }
		pagelocations.put( vpn, PL_MEMORY );
	    }
	else
	    {
		readonly = swapEntry.readOnly;
	    }

	// Record it in the inverted page table.
	TranslationEntry vmEntry = new TranslationEntry();
	vmEntry.ppn = ppn;
	vmEntry.vpn = vpn;
	vmEntry.used = false;
	vmEntry.dirty = false;
	vmEntry.valid = true;
	vmEntry.readOnly = readonly;
	ipt.write(pid, vmEntry);

	String msg = String.format("After swapping vpn %d for pid %d into " +
				   "ppn %d, in, vm is %s and swap is %s",
				   vpn, pid, ppn, ipt, swap);
	debug(dbgVM, msg);
    }

    /**
     * Save the state of this process in preparation for a context switch.
     * Called by <tt>UThread.saveState()</tt>.
     *
     * We simply set the valid bit for every entry in the TLB to
     * false.
     */
    public void saveState() {
	Processor processor = Machine.processor();
	int tlbSize = processor.getTLBSize();
	for (int i = 0; i < tlbSize; i++) {
	    TranslationEntry e = processor.readTLBEntry(i);
	    e.valid = false;
	    processor.writeTLBEntry(i, e);
	}
    }

    /**
     * The processor sets the used bit in the TranslationEntry in the
     * TLB, but I think we need to occasionally transfer that usage
     * information into our page table.
     */
    public void transferUsedBits() {
	Processor processor = Machine.processor();
	int tlbSize = processor.getTLBSize();
	for (int i = 0; i < tlbSize; i++) {
	    TranslationEntry e = processor.readTLBEntry(i);
	    if (e.valid && e.used) {
		ipt.markUsed(e.ppn);
	    }
	    if (e.valid && e.dirty) {
		ipt.markDirty(e.ppn);
	    }
	}
    }

    /**
     * Restore the state of this process after a context switch. Called by
     * <tt>UThread.restoreState()</tt>.
     */
    public void restoreState() {
	// If we call super.restoreState(), we'll get an exception
	// right away because it tries to access the Processor's page
	// table, and the Processer is using a TLB instead of a page
	// table. I don't actually think we need to do anything here
	// yet for part 1, maybe for part 2 or 3 we will.
	//	debug(dbgVM, "Restoring state of " + this);
	// super.restoreState();
    }

    /**
     * Initializes page tables for this process so that the executable can be
     * demand-paged.
     *
     * @return	<tt>true</tt> if successful.
     */
    protected boolean loadSections() {
	debug(dbgVM, "Loading sections");

	//  We'll load the program into contiguous pages of virtual
	//  memory. This counter always represents the number of the
	//  next virtual page to use.
	int vpn = 0;

	// Tom - this code should not load the page into memory anymore
	// it should wait for a page fault
	
	// load sections
	for (int s = 0; s<coff.getNumSections(); s++) {
	    CoffSection section = coff.getSection(s);
	  	    
	    Lib.debug(dbgProcess, "\tinitializing " + section.getName()
		      + " section (" + section.getLength() + " pages)");

	    // Load each page of this section
	    for (int i=0; i < section.getLength(); i++, vpn++) {

		Lib.debug(dbgProcess,
			  String.format("  loading section %d page %d at vpn %d",
					s, i, vpn));
		
		// Until we implement swapping pages out to disk, just
		// assume we can use the next page of physical memory.
		//int ppn = UserKernel.getNextPage();

		swapLock.acquire();

		/*if ( ppn == -1 ) {
		    Lib.debug( dbgProcess, "\tOut of Memory " );
		    debug(dbgVM, "Swapping out in order to load section");
		    ppn = swapOut();
		    }*/
		
		// Make a new translation entry and put it in our
		// inverted page table.
		/*TranslationEntry e = new TranslationEntry();
		e.vpn      = vpn;
		e.ppn      = ppn;
		e.valid    = true; 
		e.valid    = false; 
		e.used     = false; 
		e.dirty    = false;
		e.readOnly = section.isReadOnly();
		ipt.write(pid, e);*/

		pagelocations.put( vpn, PL_COFF );

		// Tom - this is shut off because we don't want to
		// load pages until needed
		  // Load this page of the section into physical memory.
		  //section.loadPage(i, ppn);

		// put a place in the vpnToXX tables (needed for locating the page when ready)
		vpnToS.put( vpn, s );
		vpnToSI.put( vpn, i );
		
		swapLock.release();
	    }
	}
	
	// Allocate remaining pages.
	for (; vpn < numPages; ++vpn ) {
	    
	    swapLock.acquire();

	    /*int ppn = UserKernel.getNextPage();

	    if ( ppn == -1 ) {
		Lib.debug( dbgProcess, "\tOut of Memory (2)" );
		debug(dbgVM, "Swapping out in order to allocate mem");
		ppn = swapOut();
	    }
	    
	    TranslationEntry e = new TranslationEntry();
	    e.vpn      = vpn;
	    e.ppn      = ppn;
	    e.valid    = true;
	    e.used     = true;
	    e.dirty    = false;
	    e.readOnly = false;
	    ipt.write(pid, e);*/

	    pagelocations.put( vpn, PL_ST_ARG );

	    swapLock.release();
	}
	debug(dbgVM, "Done loading sections");
	return true;
    }

    /**
     * Release any resources allocated by <tt>loadSections()</tt>.
     */
    protected void unloadSections() {
	debug(dbgVM, "Unloading sections for " + pid + ", vm is " + ipt);
	List<Integer> myPages = ipt.invalidateProcess(pid);
	swap.invalidateProcess(pid);
	for (int ppn : myPages) {
	    UserKernel.freePage(ppn);	    
	}
    }    

    /**
     * Handle a user exception. Called by
     * <tt>UserKernel.exceptionHandler()</tt>. The
     * <i>cause</i> argument identifies which exception occurred; see the
     * <tt>Processor.exceptionZZZ</tt> constants.
     *
     * @param	cause	the user exception that occurred.
     */
    public void handleException(int cause) {
	//	debug(dbgVM, "Handling exception " + cause);
	Processor processor = Machine.processor();

	switch (cause) {

	case Processor.exceptionReadOnly:
	    int a = processor.readRegister(Processor.regBadVAddr);
	    throw new RuntimeException("Error trying to write to read-only address " + a);

	case Processor.exceptionTLBMiss:
	    debug(dbgVM, "TLB Miss");
	    swapLock.acquire();

	    transferUsedBits();

	    // Get the virtual address that caused the TLB miss, and
	    // then get its page number from our page table. When we
	    // implement demand paging, we will need to handle the
	    // case when the page isn't loaded.
	    int vaddr = processor.readRegister(Processor.regBadVAddr);
	    int vpn = Processor.pageFromAddress(vaddr);
	
	    // Find the mapping in our page table, and then write an
	    // identical entry into the TLB. It's important to create
	    // a copy of the entry for the TLB, so we can change the
	    // two independently (I think).

	    // Look up the virtual to physical mapping from our page table.
	    TranslationEntry e = ensureLoaded(vpn);

	    // Copy the entry from our page table into the TLB.  TODO:
	    // Do we need to do anything if we overwrite a TLB entry?
	    // I don't think so.
	    int tlbVictim = random.nextInt(processor.getTLBSize());
	    processor.writeTLBEntry(tlbVictim, e);
	    swapLock.release();
	    return;

	case Processor.exceptionPageFault:
	    System.out.println("The Processor through a PageFault.\n");
	    terminateProcess(-1);
	    return;
	default:
	    super.handleException(cause);
	    break;
	}
    }

    /**
     * Loads the page with the given virtual page number (for the
     * current process) into physical memory if it is not already
     * loaded, and returns the TranslationEntry for it.
     */
    private TranslationEntry ensureLoaded(int vpn) {
	TranslationEntry e = ipt.getEntry(pid, vpn, false);
	numPageAccesses++;
	if ( e == null || !e.valid) {
	    numPageFaults++;
	    int ppn;
	    if ( e == null )
		ppn = VMKernel.getNextPage();
	    else
		ppn = ipt.findFreePage();

	    if (ppn == -1) {
		debug(dbgVM, "Swapping out ppn " + ppn + " due to tlb miss");
		ppn = swapOut();
	    }
	    swapIn(vpn, ppn);
	    e = ipt.table[ppn].te;

	    Lib.assertTrue(ipt.table[ppn].pid == pid);
	    Lib.assertTrue(ipt.table[ppn].te.vpn == vpn);
	}
	return e;
    }
    
    private static final int pageSize = Processor.pageSize;
    public static final char dbgProcess = 'a';
    public static final char dbgVM = 'v';
    public static final char dbgH = 'h';
    

    private static final Random random = new Random();

    private static InvertedPageTable ipt = null;
    public static Swap swap = null;

    // this is per process
    // hash of vpn to location
    private HashMap< Integer, Integer > pagelocations;
    private static final int PL_UNKNOWN = 0;
    private static final int PL_COFF    = 1;
    private static final int PL_ST_ARG  = 2;
    private static final int PL_SWAP    = 3;
    private static final int PL_MEMORY  = 4;
    // this is vpn to pair of a section and i offset in section
    // this is needed to determine the code page when needed (only occurs once)
    private HashMap< Integer, Integer> vpnToS;
    private HashMap< Integer, Integer> vpnToSI;

    private static byte[] blankpage;

    /**
     * Checks that the virtual address is valid and that offset and
     * length are valid for the given data array. If they're valid,
     * returns true, otherwise terminates the process and returns
     * false.
     */
    private boolean checkMem(int vaddr, byte[] data, int offset, int length) {
	Lib.assertTrue(offset >= 0 && 
		       length >= 0 && 
		       offset+length <= data.length);
	if (validVirtualAddress(vaddr)) {
	    return true;
	}
	// seg fault - terminate 
	terminateProcess( -1 );
	return false;
    }

    /**
     * Transfer data from this process's virtual memory to the specified array.
     * This method handles address translation details. This method must
     * <i>not</i> destroy the current process if an error occurs, but instead
     * should return the number of bytes successfully copied (or zero if no
     * data could be copied).
     *
     * @param	vaddr	the first byte of virtual memory to read.
     * @param	data	the array where the data will be stored.
     * @param	offset	the first byte to write in the array.
     * @param	length	the number of bytes to transfer from virtual memory to
     *			the array.
     * @return	the number of bytes successfully transferred.
     */
    public int readVirtualMemory(int vaddr, byte[] data, int offset,
				 int length) {
	
	if (!checkMem(vaddr, data, offset, length))
	    return -1;

	byte[] memory = Machine.processor().getMemory();
	
	int vpn = Processor.pageFromAddress(vaddr);
	int pgoff  = Processor.offsetFromAddress(vaddr);
	TranslationEntry e = ipt.getEntry(pid, vpn, true);
	int ppn = e.ppn;
	debug(dbgVM, "Reading mem at " + e.ppn);
	// Since we're using virtual memory now, vaddr may be greater
	// than memory.length. Maybe we should check to make sure that
	// the specified length doesn't reach across a page boundary.
	// int amount = Math.min(length, memory.length-vaddr);
	int amount = length;
	int realaddr = Processor.makeAddress(ppn, pgoff);
	System.arraycopy(memory, realaddr, data, offset, amount);
	e.used = true;
	//	debug(dbgVM, "After reading, vm is "/* + ipt*/);
	return amount;
    }

    /**
     * Transfer data from the specified array to this process's virtual memory.
     * This method handles address translation details. This method must
     * <i>not</i> destroy the current process if an error occurs, but instead
     * should return the number of bytes successfully copied (or zero if no
     * data could be copied).
     *
     * @param	vaddr	the first byte of virtual memory to write.
     * @param	data	the array containing the data to transfer.
     * @param	offset	the first byte to transfer from the array.
     * @param	length	the number of bytes to transfer from the array to
     *			virtual memory.
     * @return	the number of bytes successfully transferred.
     */
    public int writeVirtualMemory(int vaddr, byte[] data, int offset,
				  int length) {
	if (!checkMem(vaddr, data, offset, length))
	    return -1;

	byte[] memory = Machine.processor().getMemory();

	int vpn = Processor.pageFromAddress( vaddr );
	int pgoff  = Processor.offsetFromAddress( vaddr );
	// Since we're using virtual memory now, vaddr may be greater
	// than memory.length. Maybe we should check to make sure that
	// the specified length doesn't reach across a page boundary.
	// int amount = Math.min(length, memory.length-vaddr);
	int amount = length;
	TranslationEntry e = ensureLoaded(vpn);
	int ppn = e.ppn;

	debug(dbgVM, this + " writing to " + ppn);


	if ( e.readOnly == true ) {
	    // Terminate and return -1 to indicate segfault.
	    terminateProcess( -1 );
	    return -1;
	}

	int realaddr = Processor.makeAddress( ppn,  pgoff );
	debug(dbgVM, "Length is " + length);
	debug(dbgVM, "vaddr is " + vaddr);
	debug(dbgVM, String.format("reading offset  %d realaddr %d amount %d",
				   offset, realaddr, amount));
	System.arraycopy(data, offset, memory, realaddr, amount);

	// mark that we have written to it, this isn't doing anything
	// because we aren't swapping pages.
	e.used  = true;
	e.dirty = true; 

	return amount;
    }

}
